# Copyright (C) 2019-2022 Intel Corporation
# Copyright (C) 2024 CVAT.ai Corporation
#
# SPDX-License-Identifier: MIT

import itertools
import logging as log
import math
import os
import os.path as osp
from collections import OrderedDict, defaultdict
from functools import cached_property
from itertools import cycle
from typing import Dict, Iterable, List, Optional

import yaml

from datumaro.components.annotation import (
    Annotation,
    AnnotationType,
    Bbox,
    Points,
    PointsCategories,
    Polygon,
    Skeleton,
)
from datumaro.components.dataset_base import DatasetItem, IDataset
from datumaro.components.dataset_item_storage import ItemStatus
from datumaro.components.dataset_storage import DatasetPatch
from datumaro.components.errors import DatasetExportError, MediaTypeError
from datumaro.components.exporter import Exporter
from datumaro.components.media import Image
from datumaro.util import dump_json_file, str_to_bool
from datumaro.util.definitions import DEFAULT_SUBSET_NAME
from datumaro.util.os_util import split_path

from .format import YoloPath, YoloUltralyticsClassificationFormat, YoloUltralyticsPath


def _make_yolo_bbox(img_size, box):
    # https://github.com/pjreddie/darknet/blob/master/scripts/voc_label.py
    # <x> <y> <width> <height> - values relative to width and height of image
    # <x> <y> - are center of rectangle
    x = (box[0] + box[2]) / 2 / img_size[0]
    y = (box[1] + box[3]) / 2 / img_size[1]
    w = (box[2] - box[0]) / img_size[0]
    h = (box[3] - box[1]) / img_size[1]
    return x, y, w, h


def bbox_annotation_as_polygon(bbox: Bbox) -> List[float]:
    points = bbox.as_polygon()

    def rotate_point(x: float, y: float):
        new_x = (
            center_x
            + math.cos(rotation_radians) * (x - center_x)
            - math.sin(rotation_radians) * (y - center_y)
        )
        new_y = (
            center_y
            + math.sin(rotation_radians) * (x - center_x)
            + math.cos(rotation_radians) * (y - center_y)
        )
        return new_x, new_y

    if rotation_radians := math.radians(bbox.attributes.get("rotation", 0)):
        center_x = bbox.x + bbox.w / 2
        center_y = bbox.y + bbox.h / 2
        points = [
            coordinate
            for x, y in zip(points[::2], points[1::2])
            for coordinate in rotate_point(x, y)
        ]
    return points


def _resolve_subsets(initial_subsets: Dict[str, IDataset]) -> Dict[str, Iterable]:
    assert YoloPath.DEFAULT_SUBSET_NAME.lower() != DEFAULT_SUBSET_NAME.lower()

    subsets: Dict[str, Iterable] = {
        subset_name: subset
        for subset_name, subset in initial_subsets.items()
        if subset_name and subset_name != DEFAULT_SUBSET_NAME
    }
    for subset_name, subset in initial_subsets.items():
        if not subset_name or subset_name == DEFAULT_SUBSET_NAME:
            subset_name = YoloPath.DEFAULT_SUBSET_NAME
            try:
                subset_name = next(name for name in subsets if name.lower() == subset_name.lower())
            except StopIteration:
                pass
            subsets[subset_name] = itertools.chain(subsets.get(subset_name, []), subset)

    return subsets


class YoloExporter(Exporter):
    # https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects
    DEFAULT_IMAGE_EXT = ".jpg"
    RESERVED_CONFIG_KEYS = YoloPath.RESERVED_CONFIG_KEYS

    @classmethod
    def build_cmdline_parser(cls, **kwargs):
        parser = super().build_cmdline_parser(**kwargs)
        parser.add_argument(
            "--add-path-prefix",
            default=True,
            type=str_to_bool,
            help="Add the 'data/' prefix for paths in the dataset info (default: %(default)s)",
        )
        return parser

    def __init__(
        self, extractor: IDataset, save_dir: str, *, add_path_prefix: bool = True, **kwargs
    ) -> None:
        super().__init__(extractor, save_dir, **kwargs)

        self._prefix = "data" if add_path_prefix else ""

    def apply(self):
        save_dir = self._save_dir

        if self._extractor.media_type() and not issubclass(self._extractor.media_type(), Image):
            raise MediaTypeError("Media type is not an image")

        os.makedirs(save_dir, exist_ok=True)

        if self._save_dataset_meta:
            self._save_meta_file(self._save_dir)

        subset_lists = OrderedDict()

        subsets = _resolve_subsets(self._extractor.subsets())

        pbars = self._ctx.progress_reporter.split(len(subsets))
        for (subset_name, subset), pbar in zip(subsets.items(), pbars):
            if subset_name in self.RESERVED_CONFIG_KEYS:
                raise DatasetExportError(
                    f"Can't export '{subset_name}' subset in YOLO format, this word is reserved."
                )

            subset_image_dir = self._make_image_subset_folder(save_dir, subset_name)
            subset_anno_dir = self._make_annotation_subset_folder(save_dir, subset_name)
            os.makedirs(subset_image_dir, exist_ok=True)
            os.makedirs(subset_anno_dir, exist_ok=True)

            image_paths = OrderedDict()
            for item in pbar.iter(subset, desc=f"Exporting '{subset_name}'"):
                try:
                    image_fpath = self._export_media(item, subset_image_dir)
                    image_name = osp.relpath(image_fpath, subset_image_dir)
                    image_paths[item.id] = osp.join(
                        self._prefix, osp.relpath(subset_image_dir, save_dir), image_name
                    )

                    self._export_item_annotation(item, subset_anno_dir)

                except Exception as e:
                    self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset))

            if subset_list_name := self._make_subset_list_file(subset_name, image_paths):
                subset_lists[subset_name] = subset_list_name

        self._save_config_files(subset_lists)

    def _save_config_files(self, subset_lists: Dict[str, str]):
        extractor = self._extractor
        save_dir = self._save_dir
        label_categories = extractor.categories()[AnnotationType.label]
        label_ids = {label.name: idx for idx, label in enumerate(label_categories.items)}
        with open(osp.join(save_dir, "obj.names"), "w", encoding="utf-8") as f:
            f.writelines("%s\n" % l[0] for l in sorted(label_ids.items(), key=lambda x: x[1]))

        with open(osp.join(save_dir, "obj.data"), "w", encoding="utf-8") as f:
            f.write(f"classes = {len(label_ids)}\n")

            for subset_name, subset_list_name in subset_lists.items():
                f.write(
                    "%s = %s\n"
                    % (subset_name, osp.join(self._prefix, subset_list_name).replace("\\", "/"))
                )

            f.write("names = %s\n" % osp.join(self._prefix, "obj.names"))
            f.write("backup = backup/\n")

    def _make_subset_list_file(self, subset_name, image_paths):
        subset_list_name = f"{subset_name}{YoloPath.SUBSET_LIST_EXT}"
        subset_list_path = osp.join(self._save_dir, subset_list_name)
        if self._patch and subset_name in self._patch.updated_subsets and not image_paths:
            if osp.isfile(subset_list_path):
                os.remove(subset_list_path)
            return

        with open(subset_list_path, "w", encoding="utf-8") as f:
            f.writelines("%s\n" % s.replace("\\", "/") for s in image_paths.values())
        return subset_list_name

    def _export_media(self, item: DatasetItem, subset_img_dir: str) -> str:
        try:
            if not item.media or not (item.media.has_data or item.media.has_size):
                raise DatasetExportError(
                    "Failed to export item '%s': " "item has no image info" % item.id
                )

            image_name = self._make_image_filename(item)
            image_fpath = osp.join(subset_img_dir, image_name)

            if self._save_media:
                if item.media:
                    self._save_image(item, image_fpath)
                else:
                    log.warning("Item '%s' has no image" % item.id)

            return image_fpath

        except Exception as e:
            self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset))

    def _save_annotation_file(self, annotation_path, yolo_annotation):
        with open(annotation_path, "w", encoding="utf-8") as f:
            f.write(yolo_annotation)

    def _export_item_annotation(self, item: DatasetItem, subset_dir: str) -> None:
        try:
            height, width = item.media.size

            yolo_annotation = ""

            for bbox in item.annotations:
                annotation_line = self._make_annotation_line(width, height, bbox)
                if annotation_line:
                    yolo_annotation += annotation_line

            annotation_path = osp.join(subset_dir, f"{item.id}{YoloPath.LABELS_EXT}")
            os.makedirs(osp.dirname(annotation_path), exist_ok=True)

            self._save_annotation_file(annotation_path, yolo_annotation)

        except Exception as e:
            self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset))

    def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]:
        if not isinstance(anno, Bbox) or anno.label is None:
            return
        if anno.attributes.get("rotation", 0) % 360.0 > 0.00001:
            return

        values = _make_yolo_bbox((width, height), anno.points)
        string_values = " ".join("%.6f" % p for p in values)
        return "%s %s\n" % (self._map_labels_for_save[anno.label], string_values)

    @cached_property
    def _labels_to_save(self) -> List[int]:
        label_categories = self._extractor.categories()[AnnotationType.label]
        point_categories = self._extractor.categories().get(
            AnnotationType.points, PointsCategories.from_iterable([])
        )
        skeleton_sublabels = [
            label_categories.find(sublabel, parent=label_categories[parent_label_id].name)[0]
            for parent_label_id in point_categories.items
            for sublabel in point_categories.items[parent_label_id].labels
        ]
        return [
            label_id
            for label_id, label in enumerate(label_categories)
            if label_id not in skeleton_sublabels
        ]

    @cached_property
    def _map_labels_for_save(self) -> Dict[int, int]:
        return {label_id: index for index, label_id in enumerate(self._labels_to_save)}

    @staticmethod
    def _make_image_subset_folder(save_dir: str, subset: str) -> str:
        return osp.join(save_dir, f"obj_{subset}_data")

    @staticmethod
    def _make_annotation_subset_folder(save_dir: str, subset: str) -> str:
        return osp.join(save_dir, f"obj_{subset}_data")

    @classmethod
    def patch(cls, dataset: IDataset, patch: DatasetPatch, save_dir: str, **kwargs):
        conv = cls(dataset, save_dir=save_dir, **kwargs)
        conv._patch = patch
        conv.apply()

        for (item_id, subset), status in patch.updated_items.items():
            if status != ItemStatus.removed:
                item = patch.data.get(item_id, subset)
            else:
                item = DatasetItem(item_id, subset=subset)

            if not (status == ItemStatus.removed or not item.media):
                continue

            if subset == DEFAULT_SUBSET_NAME:
                subset = YoloPath.DEFAULT_SUBSET_NAME
            subset_image_dir = cls._make_image_subset_folder(save_dir, subset)
            subset_anno_dir = cls._make_annotation_subset_folder(save_dir, subset)

            image_path = osp.join(subset_image_dir, conv._make_image_filename(item))
            if osp.isfile(image_path):
                os.remove(image_path)

            ann_path = osp.join(subset_anno_dir, f"{item.id}{YoloPath.LABELS_EXT}")
            if osp.isfile(ann_path):
                os.remove(ann_path)

    @property
    def can_stream(self) -> bool:
        return True


class YoloUltralyticsDetectionExporter(YoloExporter):
    RESERVED_CONFIG_KEYS = YoloUltralyticsPath.RESERVED_CONFIG_KEYS

    def __init__(
        self,
        extractor: IDataset,
        save_dir: str,
        *,
        add_path_prefix: bool = True,
        config_file: Optional[str] = None,
        write_track_id: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(extractor, save_dir, add_path_prefix=add_path_prefix, **kwargs)
        self._write_track_id = write_track_id
        self._config_filename = config_file or YoloUltralyticsPath.DEFAULT_CONFIG_FILE

    def _save_annotation_file(self, annotation_path, yolo_annotation):
        if yolo_annotation:
            super()._save_annotation_file(annotation_path, yolo_annotation)

    @classmethod
    def build_cmdline_parser(cls, **kwargs):
        parser = super().build_cmdline_parser(**kwargs)
        parser.add_argument(
            "--config-file",
            default=YoloUltralyticsPath.DEFAULT_CONFIG_FILE,
            type=str,
            help="config file name (default: %(default)s)",
        )
        parser.add_argument(
            "--write-track-id",
            default=False,
            type=str_to_bool,
            help="save track id to annotations",
        )
        return parser

    def _save_config_files(self, subset_lists: Dict[str, str], **extra_config_fields):
        extractor = self._extractor
        save_dir = self._save_dir
        with open(osp.join(save_dir, self._config_filename), "w", encoding="utf-8") as f:
            label_categories = extractor.categories()[AnnotationType.label]
            data = dict(
                path=".",
                names={
                    index: label_categories[label_id].name
                    for label_id, index in self._map_labels_for_save.items()
                },
                **subset_lists,
                **extra_config_fields,
            )
            yaml.dump(data, f)

    @staticmethod
    def _make_image_subset_folder(save_dir: str, subset: str) -> str:
        return osp.join(save_dir, YoloUltralyticsPath.IMAGES_FOLDER_NAME, subset)

    @staticmethod
    def _make_annotation_subset_folder(save_dir: str, subset: str) -> str:
        return osp.join(save_dir, YoloUltralyticsPath.LABELS_FOLDER_NAME, subset)

    def _make_track_id_suffix(self, anno: Annotation) -> str:
        track_id = anno.attributes.get("track_id") if self._write_track_id else None
        try:
            return f" {int(track_id)}"
        except (ValueError, TypeError):
            return ""

    def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]:
        anno_line = super()._make_annotation_line(width=width, height=height, anno=anno)

        if anno_line and (track_id_suffix := self._make_track_id_suffix(anno)):
            anno_line = f"{anno_line.strip()}{track_id_suffix}\n"

        return anno_line


class YoloUltralyticsSegmentationExporter(YoloUltralyticsDetectionExporter):
    def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]:
        if anno.label is None or not isinstance(anno, Polygon):
            return
        values = [value / size for value, size in zip(anno.points, cycle((width, height)))]
        string_values = " ".join("%.6f" % p for p in values)
        return (
            f"{self._map_labels_for_save[anno.label]} "
            f"{string_values}{self._make_track_id_suffix(anno)}\n"
        )


class YoloUltralyticsOrientedBoxesExporter(YoloUltralyticsDetectionExporter):
    def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]:
        if anno.label is None or not isinstance(anno, Bbox):
            return
        points = bbox_annotation_as_polygon(anno)
        values = [value / size for value, size in zip(points, cycle((width, height)))]
        string_values = " ".join("%.6f" % p for p in values)
        return (
            f"{self._map_labels_for_save[anno.label]} "
            f"{string_values}{self._make_track_id_suffix(anno)}\n"
        )


class YoloUltralyticsPoseExporter(YoloUltralyticsDetectionExporter):
    @cached_property
    def _labels_to_save(self) -> List[int]:
        point_categories = self._extractor.categories().get(
            AnnotationType.points, PointsCategories.from_iterable([])
        )
        return sorted(point_categories.items)

    @cached_property
    def _max_number_of_points(self):
        point_categories = self._extractor.categories().get(AnnotationType.points)
        if point_categories is None or len(point_categories) == 0:
            return 0
        return max(len(category.labels) for category in point_categories.items.values())

    def _save_config_files(self, subset_lists: Dict[str, str], **extra_config_fields):
        super()._save_config_files(
            subset_lists=subset_lists,
            kpt_shape=[self._max_number_of_points, 3],
            **extra_config_fields,
        )

    def _make_annotation_line(self, width: int, height: int, skeleton: Annotation) -> Optional[str]:
        if skeleton.label is None or not isinstance(skeleton, Skeleton):
            return

        x, y, w, h = skeleton.get_bbox()
        bbox_values = _make_yolo_bbox((width, height), [x, y, x + w, y + h])
        bbox_string_values = " ".join("%.6f" % p for p in bbox_values)

        point_label_ids = [
            self._extractor.categories()[AnnotationType.label].find(
                name=child_label,
                parent=self._extractor.categories()[AnnotationType.label][skeleton.label].name,
            )[0]
            for child_label in self._extractor.categories()[AnnotationType.points]
            .items[skeleton.label]
            .labels
        ]

        points_values = [f"0.0 0.0 {Points.Visibility.absent.value}"] * self._max_number_of_points
        for element in skeleton.elements:
            assert len(element.points) == 2 and len(element.visibility) == 1
            position = point_label_ids.index(element.label)
            x = element.points[0] / width
            y = element.points[1] / height
            points_values[position] = f"{x:.6f} {y:.6f} {element.visibility[0].value}"

        return (
            f"{self._map_labels_for_save[skeleton.label]} {bbox_string_values} "
            f"{' '.join(points_values)}{self._make_track_id_suffix(skeleton)}\n"
        )


class YoloUltralyticsClassificationExporter(Exporter):
    DEFAULT_IMAGE_EXT = ".jpg"

    def apply(self):
        save_dir = self._save_dir

        if self._extractor.media_type() and not issubclass(self._extractor.media_type(), Image):
            raise MediaTypeError("Media type is not an image")

        os.makedirs(save_dir, exist_ok=True)
        self._used_paths = defaultdict(set)

        if self._save_dataset_meta:
            self._save_meta_file(self._save_dir)

        labels = self._extractor.categories()[AnnotationType.label]

        subsets = _resolve_subsets(self._extractor.subsets())
        pbars = self._ctx.progress_reporter.split(len(subsets))
        for (subset_name, subset), pbar in zip(subsets.items(), pbars):
            os.makedirs(osp.join(self._save_dir, subset_name), exist_ok=True)

            items_info = defaultdict(dict)

            for item in pbar.iter(subset, desc=f"Exporting '{subset_name}'"):
                try:
                    items_info[item.id]["labels"] = [
                        labels[anno.label].name
                        for anno in item.annotations
                        if anno.type == AnnotationType.label
                    ]
                    for label_name in items_info[item.id]["labels"] or [
                        YoloUltralyticsClassificationFormat.IMAGE_DIR_NO_LABEL
                    ]:
                        items_info[item.id]["path"] = self._export_media_for_label(
                            item, subset_name, label_name
                        )

                except Exception as e:
                    self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset))

            labels_path = osp.join(
                self._save_dir,
                subset_name,
                YoloUltralyticsClassificationFormat.LABELS_FILE,
            )
            dump_json_file(labels_path, items_info)

    def _generate_path_for_item(self, item: DatasetItem, label_folder_path: str) -> str:
        item_name = base_item_name = split_path(item.id)[-1]
        if item_name in self._used_paths[label_folder_path]:
            for index in itertools.count():
                item_name = f"{base_item_name}.{index}"
                if item_name not in self._used_paths[label_folder_path]:
                    break
        self._used_paths[label_folder_path].add(item_name)
        return self._make_image_filename(item, name=item_name, subdir=label_folder_path)

    def _export_media_for_label(self, item: DatasetItem, subset_name: str, label_name: str) -> str:
        try:
            if not item.media or not (item.media.has_data or item.media.has_size):
                raise DatasetExportError(
                    "Failed to export item '%s': " "item has no image info" % item.id
                )
            subset_path = osp.join(self._save_dir, subset_name)
            label_folder_path = osp.join(subset_path, label_name)

            image_fpath = self._generate_path_for_item(item, label_folder_path)

            if self._save_media:
                if item.media:
                    os.makedirs(label_folder_path, exist_ok=True)
                    self._save_image(item, image_fpath)
                else:
                    log.warning("Item '%s' has no image" % item.id)

            return osp.relpath(image_fpath, subset_path)
        except Exception as e:
            self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset))

    @property
    def can_stream(self) -> bool:
        return True
